import re

import traceback
import arrow

from pymongo import MongoClient, UpdateOne
from conf.settings import CLOUDMONITOR_MONGODB
from utility.log import Logger
from utility.utils import dis_number, trans_interval, backup, rollback

log = Logger()

db = MongoClient(CLOUDMONITOR_MONGODB)['cloudmonitor-dev2']
# db = MongoClient(host='localhost', port=27017)['cloudmonitor']
now_datetime = arrow.utcnow()
reset_day = now_datetime.replace(day=1, hour=2, minute=0, second=0, microsecond=0).shift(months=1, hours=-8).int_timestamp
now = now_datetime.isoformat()

task_type_tuple = ('vul', 'ssl', 'asset', 'securityEvent', 'ping', 'http', 'content')
task_type_patterns = f"^{'|'.join(task_type_tuple)}+?"
feature_tuple = (
    'httpTask', 'pingTask', 'vulTask', 'vulManualTask', 'sslTask',
    'sslManualTask', 'securityEventTask', 'securityEventManualTask',
    'assetTask', 'assetManualTask', 'smsAlert', 'emailAlert', 'reportDownload',
    'contentTask', 'contentManualTask'
)

item = {
    "used": 0,
    "total": 200000,
    "number": 0,
    "limitType": "period",
    "time": "2629800000",
    "second": 2629800,
}


def create_charge_template(balance, total, func_type, task_type):
    return {
        'updateTime': now_datetime.datetime,
        'createTime': now_datetime.datetime,
        'isDelete': False,
        'chargeType': 'plan',
        'uid': balance['uid'],
        'packId': balance['planId'],
        'limitType': 'frequency',
        'expireTime': balance.get('planChargeTime'),
        'number': 1,
        'used': 0,
        'total': total,
        'funcType': func_type,
        'taskType': task_type,
    }


def merge_feature(f0, f1):
    if len(f0) < len(f1):
        temp = f0
        f0 = f1
        f1 = temp

    new_feature = {}
    for k, v in f0.items():
        new_feature[k] = v or f1.get(k, False)

    return new_feature


def calculate_pack_feature(package):
    feature = {k: False for k in feature_tuple}
    task_feature = {}
    limit = package.get('limit', {})
    for k, v in limit.items():
        if not v.get('total'):
            continue

        if res := re.match(task_type_patterns, k):
            _task_type = res.group()
            if "Manual" in k:
                task_feature[f"{_task_type}ManualTask"] = True
                if v.get('limitType') == 'period':
                    task_feature[f"{_task_type}Task"] = True
            elif "MaxFreq" in k or "MaxTask" in k:
                task_feature[f"{_task_type}Task"] = True
        elif k not in feature_tuple:
            continue
        else:
            feature[k] = True

    feature.update(task_feature)

    feature['reportDownload'] = False

    return feature


def calculate_feature(uid: str):
    pipeline = [
        {'$match': {'uid': uid, 'isDelete': False}},
        {'$group': {'_id': '$packId'}}
    ]

    id_list = [item['_id'] for item in db['charges'].aggregate(pipeline)]

    feature = {}
    for item in db['packages'].find({'_id': {'$in': id_list}}):
        if not feature:
            feature = item['feature']
        else:
            feature = merge_feature(feature, item['feature'])

    return feature


def get_manual_total(manual_task):
    order = ('vul', 'securityEvent', 'ssl', 'asset')
    i = len(order)

    used_list = dis_number(manual_task.get('used', 0), i)
    total_list = dis_number(manual_task.get('total', 0), i)

    results = {
        f'{task_type}ManualTask': {
            'used': int(used_list[i]),
            'total': int(total_list[i]),
            'number': int(manual_task.get('number', 0)),
            'limitType': manual_task.get('limitType', 'period'),
            'time': reset_day,
            'taskType': task_type
        } for i, task_type in enumerate(order)}

    return results


def upgrade_balances(pack_cache, default_oid):
    db.drop_collection('charges')
    i = 0
    for balance in db['balances'].find():
        _new_charges = []
        extra_balance = {}

        balance['balance']['urlBLSMaxNum'] = {
            "used": db['url_blacklist_setting'].count_documents({'uid': balance['uid']}),
            "total": 50,
            "number": 1,
            "limitType": "limit",
            "taskType": "",
            "time": 0,
        }

        for k, v in balance['balance'].items():
            # freq: total = time
            for _task_type in task_type_tuple:
                if k == f"{_task_type}MaxFreq":
                    v['total'] = int(v.get('time', 0))
                    break

            # manualTask -> {taskType}manualTask
            if k == 'manualTask' and (manual_task := balance['balance'].get('manualTask')):
                for _k, _v in get_manual_total(manual_task).items():
                    if _task_type := re.match(f'({task_type_patterns})ManualTask', _k):
                        _task_type = _task_type.group(0)
                    else:
                        _task_type = ''
                    _new_charges.append(create_charge_template(balance, _v.get('total', 0), _k, _task_type))
                    extra_balance[_k] = _v

            # aviliableMaxNode -> httpMaxNode pingMaxNode
            if k == 'aviliableMaxNode' and (max_node := balance.pop('aviliableMaxNode', None)):
                extra_balance['httpMaxNode'] = max_node
                extra_balance['pingMaxNode'] = max_node
                _new_charges.append(create_charge_template(balance, max_node.get('total', 0), 'httpMaxNode', 'http'))
                _new_charges.append(create_charge_template(balance, max_node.get('total', 0), 'pingMaxNode', 'ping'))

            # set default package
            if not balance.get('planId'):
                balance['planId'] = default_oid

            # set feature
            if balance.get('planId') and (plan := pack_cache.get(balance['planId'])):
                balance['feature'] = plan['feature']

            task_type = ''
            for _task_type in task_type_tuple:
                if _task_type in k:
                    task_type = _task_type
                    break
            v['taskType'] = task_type
            # add charge
            if k in ('aviliableMaxNode', 'manualTask'):
                continue

            for _charges in v.pop('charges', []):
                if not _charges.get('expireTime'):
                    _charges['expireTime'] = balance['planExpireTime']

                _charges.pop('_id', None)
                _charges['uid'] = balance['uid']
                _charges['funcType'] = k
                _charges['packId'] = _charges.pop('ref', None)
                _charges['taskType'] = task_type
                _charges['isDelete'] = False
                if not _charges.get('createTime'):
                    _charges['createTime'] = now_datetime.datetime
                _charges['updateTime'] = now_datetime.datetime

                if _charges['limitType'] == 'period':
                    _charges['time'] = reset_day
                else:
                    _charges['time'] = 0

                _new_charges.append(_charges)

        # add extra balance
        balance['balance'].update(extra_balance)

        # del item
        balance['balance'].pop('manualTask', None)

        # extra freq
        _temp_dict = create_freq(balance['balance'])
        _temp_dict['contentManualTask'] = {
            "used": 0,
            "total": balance['balance'].get('domain', {}).get('total', 0),
            "number": 1,
            "limitType": "period",
            "taskType": 'content',
            "time": reset_day
        }
        for _k, _v in _temp_dict.items():
            balance['balance'][_k] = _v
            _new_charges.append(create_charge_template(balance, _v.get("total", 0), _k, _v.get('taskType')))

        db['balances'].update_one({'_id': balance['_id']}, {'$set': balance})
        log.logger.info(f"Updated balances. uid {balance['uid']}")
        i += 1

        if _new_charges:
            db['charges'].insert_many(_new_charges)
            calculate_feature(balance.get('uid'))
            log.logger.info(f"Created charges for uid {balance['uid']}")

    return i


def upgrade_jobs():
    i = 0

    update_list = []
    for job in db['jobs'].find():
        for _task_type in task_type_tuple:
            if job.get(f'{_task_type}Settings'):
                _settings = job[f'{_task_type}Settings']
                if isinstance(_settings, list):
                    for _setting in _settings:
                        _setting['interval'] = trans_interval(_setting['interval'])
                else:
                    _settings['interval'] = trans_interval(_settings['interval'])
        update_list.append(UpdateOne({'_id': job['_id']}, {'$set': job}))
        log.logger.info(f"Updated job. id: {job['_id']}")
        i += 1

    if update_list:
        db['jobs'].bulk_write(update_list)

    update_list = []
    for asset in db['assets'].find():
        _setting = asset['taskSettings']
        _setting['interval'] = trans_interval(_setting['interval'])
        update_list.append(UpdateOne({'_id': asset['_id']}, {'$set': asset}))
        log.logger.info(f"Updated asset. id: {asset['_id']}")
        i += 1

    if update_list:
        db['assets'].bulk_write(update_list)

    return i


def create_default_package():
    plan = {
        'createTime': arrow.utcnow().datetime,
        'key': 'scanv_free',
        'name': '免费版',
        'marks': '免费版',
        'operator': 'scanv',
        'isDefaultPlan': True,
        'chargeType': 'plan',
        'feature': {},
        'limit': {},
        'packageMarks': '',
    }

    insert_id = db['packages'].insert_one(plan).inserted_id

    log.logger.info('Succeed create default plan.')

    return insert_id


def create_packages():
    # condition: additionPackages._id != plan._id

    db.drop_collection('packages')

    insert_list = []
    pack_cache = {}
    i = 0
    for plan in db['plans'].find():
        plan['marks'] = '旧版套餐'
        plan['operator'] = 'scanv'
        plan['chargeType'] = 'plan'
        plan['createTime'] = now_datetime.datetime

        # manual task
        plan['limit'].update(get_manual_total(plan['limit'].pop('manualTask')))
        plan['limit']['contentManualTask'] = {
            "used": 0,
            "total": plan['limit'].get('domain', {}).get('total', 0),
            "number": 1,
            "limitType": "period",
            "taskType": 'content',
            "time": reset_day
        }

        # aviliableMaxNode -> httpMaxNode pingMaxNode
        max_node = plan['limit'].pop('aviliableMaxNode')
        plan['limit']['httpMaxNode'] = max_node
        plan['limit']['pingMaxNode'] = max_node

        plan['isDefaultPlan'] = False

        for func_type, v in plan['limit'].items():
            if (v['limitType'] == 'frequency') and v.get('time'):
                v['total'] = int(v['time'])
                v['time'] = 0

        plan['limit'].update(create_freq(plan['limit']))
        plan['limit']['urlBLSMaxNum'] = {
            "used": 0,
            "total": 50,
            "number": 1,
            "limitType": "limit",
            "taskType": "",
            "time": 0,
        }

        # feature
        plan['feature'] = calculate_pack_feature(plan)
        plan['feature']['reportDownload'] = True

        insert_list.append(plan)
        pack_cache[plan['_id']] = plan
        log.logger.info(f"Create packages. id: {plan['_id']}")
        i += 1

    for addition in db['additionpackages'].find():
        addition['marks'] = '旧版模块'
        addition['operator'] = 'scanv'
        addition['chargeType'] = 'additionPackage'
        addition['feature'] = {}
        addition['packageMarks'] = '旧版套餐'
        addition['isDefaultPlan'] = False
        addition['createTime'] = now_datetime.datetime
        _old_limit = addition.pop('package')
        addition['limit'] = _old_limit

        # manual task
        addition['limit'].update(get_manual_total(_old_limit['manualTask']))
        addition['limit']['contentManualTask'] = {
            "used": 0,
            "total": addition['limit'].get('domain', {}).get('total', 0),
            "number": 1,
            "limitType": "period",
            "taskType": 'content',
            "time": reset_day
        }

        # aviliableMaxNode -> httpMaxNode pingMaxNode
        max_node = addition['limit'].pop('aviliableMaxNode')
        addition['limit']['httpMaxNode'] = max_node
        addition['limit']['pingMaxNode'] = max_node

        for func_type, v in addition['limit'].items():
            if (v['limitType'] == 'frequency') and v.get('time'):
                v['total'] = int(v['time'])
                v['time'] = 0

        addition['limit'].update(create_freq(addition['limit']))

        addition['feature'] = calculate_pack_feature(addition)
        insert_list.append(addition)
        pack_cache[addition['_id']] = addition
        log.logger.info(f"Create packages. id: {addition['_id']}")
        i += 1

    if insert_list:
        db['packages'].insert_many(insert_list)

    return i, pack_cache


def create_freq(limit):
    # content
    temp_dict = {}
    use_total_security = False
    use_total_vul = False

    if limit.get('securityEventManualTask', {}).get('total'):
        use_total_security = True
    if limit.get('vulManualTask', {}).get('total') or limit.get('vulMaxTask', {}).get('total'):
        use_total_vul = True

    temp_dict.update({
        'securityEventMaxPage': {
            "used": 0,
            "total": 100000 if use_total_security else 0,
            "number": 1,
            "limitType": "frequency",
            "taskType": "securityEvent",
            "time": 0,
        },
        'contentMaxPage': {
            "used": 0,
            "total": 100000 if use_total_security else 0,
            "number": 1,
            "limitType": "frequency",
            "taskType": "content",
            "time": 0,
        },
        'contentMaxFreq': {
            "used": 0,
            "total": 86400000 if use_total_security else 0,
            "number": 1,
            "limitType": "frequency",
            "taskType": "content",
            "time": 0,
        },
        'contentMaxDepth': {
            "used": 0,
            "total": 10 if use_total_security else 0,
            "number": 1,
            "limitType": "frequency",
            "taskType": "content"
        },
        "securityEventMaxDepth": {
            "used": 0,
            "total": 10 if use_total_security else 0,
            "number": 1,
            "limitType": "frequency",
            "taskType": "securityEvent",
            "time": 0,
        },
    })

    temp_dict.update({
        'vulMaxPage': {
            "used": 0,
            "total": 100000 if use_total_vul else 0,
            "number": 1,
            "limitType": "frequency",
            "taskType": "vul",
            "time": 0,
        },
        'vulMaxDepth': {
            "used": 0,
            "total": 10 if use_total_vul else 0,
            "number": 1,
            "limitType": "frequency",
            "taskType": "vul",
            "time": 0,

        },
    })

    return temp_dict


def main():
    backup_collections = ('balances', 'jobs', 'assets')
    backup(backup_collections, log.logger, db, now)
    try:
        pack_count, pack_cache = create_packages()
        default_oid = create_default_package()
        balances_count = upgrade_balances(pack_cache, default_oid)
        job_count = upgrade_jobs()

        log.logger.info(f"Result: \n"
                    f"Created {pack_count} packages.\n"
                    f"Upgraded {balances_count} balances.\n"
                    f"Upgraded {job_count} jobs/assets.\n")
    except Exception as e:
        log.logger.error(e)
        traceback.print_exc()
    # finally:
        rollback(backup_collections, log.logger, db, now)


if __name__ == '__main__':
    main()

    # now = '2021-09-24T02:07:44.808493+00:00'
    # rollback(('balances', 'jobs', 'assets'))
